import cv2
import numpy as np
import tensorrt as trt
import pycuda.autoinit
import pycuda.driver as cuda
import time
import ctypes
from threading import Thread


class YOLOTensorRTDetector:
    def __init__(self, engine_path):
        # 加载 TensorRT 引擎
        logger = trt.Logger(trt.Logger.INFO)
        runtime = trt.Runtime(logger)

        with open(engine_path, 'rb') as f:
            engine = runtime.deserialize_cuda_engine(f.read())

        self.context = engine.create_execution_context()
        self.stream = cuda.Stream()

        # 设置输入输出
        self.input_shape = (1, 3, 640, 640)
        self.output_shape = (1, 6001)

        # 分配内存
        host_memory = cuda.pagelocked_empty(trt.volume(self.input_shape), dtype=np.float32)
        device_memory = cuda.mem_alloc(host_memory.nbytes)

        # 存储缓冲区信息
        self.bindings = [int(device_memory)]
        self.host_inputs = [host_memory]
        self.cuda_inputs = [device_memory]

        # 获取输入张量名称
        self.input_names = []
        for i in range(engine.num_io_tensors):
            name = engine.get_tensor_name(i)
            if engine.get_tensor_mode(name) == trt.TensorIOMode.INPUT:
                self.input_names.append(name)
                print(f"Input tensor name: {name}")

    def preprocess(self, image):
        # 图像预处理
        resized = cv2.resize(image, (640, 640))
        rgb = cv2.cvtColor(resized, cv2.COLOR_BGR2RGB)
        normalized = rgb.astype(np.float32) / 255.0
        transposed = normalized.transpose(2, 0, 1)
        return np.expand_dims(transposed, axis=0)

    def infer(self, image):
        # 预处理
        input_data = self.preprocess(image)
        np.copyto(self.host_inputs[0], input_data.ravel())

        # 拷贝到 GPU
        cuda.memcpy_htod_async(self.cuda_inputs[0], self.host_inputs[0], self.stream)

        # 执行推理
        self.context.execute_async_v3(stream_handle=self.stream.handle)

        # 同步流
        self.stream.synchronize()

        return True


def test_single_camera_dual_exposure():
    CAMERA_INDEX = 0
    EXPOSURE_1 = 100  # 曝光值1
    EXPOSURE_2 = -1  # 曝光值2

    cap = cv2.VideoCapture(CAMERA_INDEX)
    if not cap.isOpened():
        print(f"无法打开摄像头 {CAMERA_INDEX}")
        return

    # 先设置一次曝光，避免部分摄像头第一次设置无效
    cap.set(cv2.CAP_PROP_AUTO_EXPOSURE, EXPOSURE_1)
    time.sleep(0.2)

    detector = YOLOTensorRTDetector("build/yolov5s.engine")


    try:
        while True:
            # 设置曝光1，采集一帧
            cap.set(cv2.CAP_PROP_AUTO_EXPOSURE, EXPOSURE_1)
            time.sleep(0.1)  # 等待曝光生效
            ret1, frame1 = cap.read()
            if not ret1:
                print("无法获取画面1")
                break

            # 执行推理
            start_time = time.time()
            detector.infer(frame1)
            fps = 1.0 / (time.time() - start_time)

            # 显示FPS
            cv2.putText(frame1, f"FPS: {fps:.1f}", (10, 30),
                        cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
            cv2.putText(frame1, "Exposure -3", (10, 60),
                        cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)

            # 显示结果
            cv2.imshow("Camera1", frame1)

            if cv2.waitKey(1) & 0xFF == ord('q'):
                break

            # 设置曝光2，采集一帧
            cap.set(cv2.CAP_PROP_AUTO_EXPOSURE, EXPOSURE_2)
            time.sleep(0.1)
            ret2, frame2 = cap.read()
            if not ret2:
                print("无法获取画面2")
                break

            # 执行推理
            start_time = time.time()
            detector.infer(frame2)
            fps = 1.0 / (time.time() - start_time)

            # 显示FPS
            cv2.putText(frame2, f"FPS: {fps:.1f}", (10, 30),
                        cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
            cv2.putText(frame2, "Exposure -1", (10, 60),
                        cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)

            # 显示结果
            cv2.imshow("Camera2", frame2)

            if cv2.waitKey(1) & 0xFF == ord('q'):
                break
    finally:
        cap.release()
        cv2.destroyAllWindows()


if __name__ == "__main__":
    test_single_camera_dual_exposure()
